from openai import OpenAI

class VLLMClient:
    def __init__(self, model_name):
        openai_api_key = "EMPTY"
        openai_api_base = "http://localhost:6006/v1"
        
        self.client = OpenAI(
            api_key=openai_api_key,
             base_url=openai_api_base,
        )
        self.model_name = model_name
        
    def get_response(self, text, temperature=0.7, top_p=0.9, max_tokens=1024):
        messages = [{"role": "user", "content": text}]
        
        completion = self.client.chat.completions.create(
            model=self.model_name,
            messages = messages,
            temperature=temperature,
            top_p=top_p,
            max_tokens=max_tokens
        )
        
        return completion.choices[0].message.content
    
if __name__ == '__main__':
    client = VLLMClient("Qwen2.5-14B-Instruct-1M")
    
    text = '''
# 将下列用户提问拆分为多个子问题,当查询完所有的子问题时，即可基于每个子问题的答案得到原问题的答案。要求：
## 1. 只列出子问题即可，不要输出其他内容。
## 2. 每行输出一步查询的子问题，后续的子问题可能依赖前边子问题的答案。

# 用户提问：
比较一下中海达卫星导航和长远锂科2020年的净利润。
    '''
    
    response = client.get_response(text, temperature=0, top_p=0.8, max_tokens=1024)
    print(response)